package com.gitee.wsl.io.channel.write.ext

import com.gitee.wsl.io.channel.ByteChannel
import com.gitee.wsl.io.channel.ChannelJob
import com.gitee.wsl.io.channel.read.ByteReadChannel
import com.gitee.wsl.io.channel.write.ByteWriteChannel
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.InternalCoroutinesApi
import kotlinx.coroutines.Job
import kotlinx.coroutines.cancel
import kotlinx.coroutines.job
import kotlinx.coroutines.launch
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.EmptyCoroutineContext


class WriterScope(
    val channel: ByteWriteChannel,
    override val coroutineContext: CoroutineContext
) : CoroutineScope

class WriterJob internal constructor(
    val channel: ByteReadChannel,
    override val job: Job
) : ChannelJob

@Suppress("UNUSED_PARAMETER")
fun CoroutineScope.writer(
    coroutineContext: CoroutineContext = EmptyCoroutineContext,
    autoFlush: Boolean = false,
    block: suspend WriterScope.() -> Unit
): WriterJob = writer(coroutineContext, ByteChannel(), block)

@OptIn(InternalCoroutinesApi::class)
fun CoroutineScope.writer(
    coroutineContext: CoroutineContext = EmptyCoroutineContext,
    channel: ByteChannel,
    block: suspend WriterScope.() -> Unit
): WriterJob {
    val job = launch(coroutineContext) {
        val nested = Job(this.coroutineContext.job)
        try {
            block(WriterScope(channel, this.coroutineContext + nested))
            nested.complete()

            if (this.coroutineContext.job.isCancelled) {
                channel.cancel(this.coroutineContext.job.getCancellationException())
            }
        } catch (cause: Throwable) {
            nested.cancel("Exception thrown while writing to channel", cause)
            channel.cancel(cause)
        } finally {
            nested.join()
            runCatching { channel.flushAndClose() }
        }
    }.apply {
        invokeOnCompletion {
            if (it != null && !channel.isClosedForWrite) {
                channel.cancel(it)
            }
        }
    }

    return WriterJob(channel, job)
}